cmake_minimum_required(VERSION 3.10)
project(TrackNetV3_BallDetection)

set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

enable_language(CUDA)

# 查找必要的包
find_package(OpenCV REQUIRED)
find_package(PkgConfig REQUIRED)

# 查找FFmpeg
pkg_check_modules(FFmpeg REQUIRED libavcodec libavformat libavutil libswscale)

# 设置TensorRT路径（根据实际安装路径修改）
set(TENSORRT_DIR "/home/imvision/software/TensorRT-8.6.1.6" CACHE PATH "TensorRT installation directory")
find_path(TENSORRT_INCLUDE_DIR NvInfer.h
    PATHS ${TENSORRT_DIR}/include
    NO_DEFAULT_PATH)
find_library(TENSORRT_LIBRARY nvinfer
    PATHS ${TENSORRT_DIR}/lib
    NO_DEFAULT_PATH)
find_library(TENSORRT_PARSER_LIBRARY nvonnxparser
    PATHS ${TENSORRT_DIR}/lib
    NO_DEFAULT_PATH)

# CUDA
find_package(CUDAToolkit REQUIRED)

# 包含目录
include_directories(
    ${CMAKE_CURRENT_SOURCE_DIR}
    ${TENSORRT_INCLUDE_DIR}
    ${FFmpeg_INCLUDE_DIRS}
    ${OpenCV_INCLUDE_DIRS}
)

set(CUDA_SOURCES
    yuv2rgb.cu
)
# 可执行文件
add_executable(tracknetv3_demo
    main.cpp
    tracknetv3_trt.cpp
    video_processor.cpp
    ball_tracker.cpp
    ${CUDA_SOURCES}
)

set_target_properties(tracknetv3_demo PROPERTIES
    CUDA_ARCHITECTURES "75"  # 例如 Turing 架构，根据你的 GPU 调整
)

# 链接库
target_link_libraries(tracknetv3_demo
    ${OpenCV_LIBS}
    ${FFmpeg_LIBRARIES}
    avcodec
    avformat
    avutil
    swscale
    ${TENSORRT_LIBRARY}
    ${TENSORRT_PARSER_LIBRARY}
    CUDA::cudart
    pthread
    dl
)


# 编译器选项
target_compile_options(tracknetv3_demo PRIVATE -Wall -O2)
target_compile_definitions(tracknetv3_demo PRIVATE 
    $<$<PLATFORM_ID:Linux>:__linux__>
    USE_GPU_DECODER=1
)

# 安装目标（可选）
install(TARGETS tracknetv3_demo DESTINATION bin)
